#!/usr/bin/python
# -*- coding: utf-8 -*-
from torch import Tensor

"""
TP（True Positive）：正确的正例，一个实例是正类并且也被判定成正类
FN（False Negative）：错误的反例，漏报，本为正类但判定为假类
FP（False Positive）：错误的正例，误报，本为假类但判定为正类
TN（True Negative）：正确的反例，一个实例是假类并且也被判定成假类
准确率：所有预测正确（正类负类）的占总比重：Accuracy = （TP + TN）/(TP + TN + FP + FN)
精确率：（也叫查准率）正确预测为正的占全部预测为正的比例（不准错，宁愿漏检，也不能让现有的预测有错） Precision = TP / (TP + FP)
召回率：即正确预测为正的占全部实际为正的比例（不准漏，宁可错杀一千，也不放过一个）Recall = TP / (TP + FN)
F1分数（F1-score）是分类问题的一个衡量指标。一些多分类问题的机器学习竞赛，常常将F1-score作为最终测评的方法。它是精确率和召回率的调和平均数，最大为1，最小为0。
    F1-score = 2 * ((Precision * Recall)/(Precision + Recall))
"""


class MultiClassEvaluation:
    def __init__(self, category_names:list, detailed_str=False, keep_data=True):
        self.detailed_str = detailed_str
        self.category_names = category_names
        self.category_len = len(category_names)
        self.keep_data = keep_data
        self.data:list = list()
        self.evaluations = dict()
        self.total_eval = ModelEvaluation(detailed_str=detailed_str)
        for i in range(self.category_len):
            self.evaluations[self.category_names[i]] = ModelEvaluation(detailed_str=detailed_str)

    def get_evaluation(self, index):
        return self.evaluations[self.category_names[index]]

    def add(self, real:Tensor, output:Tensor, loss=0, clear=True):
        """

        :param real:
        :param output:
        :param loss:
        :param clear: 是否清空指标数据
        :return:
        """
        def cell_add(real_cell, output_cell):
            """
            多字段多分类
            :param real_cell:
            :param output_cell:
            :return:
            """
            self.save_data(real_cell, output_cell, loss=loss)
            top_r_n, top_r_i = real_cell.topk(1)
            top_r_i = top_r_i[0].item()
            top_o_n, top_o_i = output_cell.topk(1)
            top_o_i = top_o_i[0].item()
            for i in range(self.category_len): # 分量指标计算
                e = self.get_evaluation(i)
                if i == top_r_i == top_o_i:
                    e.add_tp()
                elif i == top_r_i and i != top_o_i:
                    e.add_fn()
                elif i != top_r_i and i == top_o_i:
                    e.add_fp()
                elif i != top_r_i and i != top_o_i:
                    e.add_tn()

        def cell_one_label_add(real_cell, output_cell):
            """
            一个字段的二分类
            :param real_cell:
            :param output_cell:
            :return:
            """
            e = self.get_evaluation(0)
            if real_cell[0] == 1 and 1 - output_cell[0] < 0.5:
                e.add_tp()
            elif real_cell[0] == 1 and 1 - output_cell[0] >= 0.5:
                e.add_fn()
            elif real_cell[0] == 0 and output_cell[0] > 0.5:
                e.add_fp()
            elif real_cell[0] == 0 and output_cell[0] <= 0.5:
                e.add_tn()
        if clear:
            for i in range(self.category_len):
                self.evaluations[self.category_names[i]].clear()
        if len(real.shape) == 2 and len(output.shape) == 2:
            for i in range(real.shape[0]):
                if real.shape[1] == 1:
                    cell_one_label_add(real[i], output[i])
                else:
                    cell_add(real[i], output[i])
            return True
        elif len(real.shape) == 1:
            if len(output.shape) == 2:
                output = output.squeeze(0)
            if len(real) == 1:
                cell_one_label_add(real, output)
            else:
                cell_add(real, output)
            return True
        return False

    def save_data(self, real, output, loss=0):
        if self.keep_data:
            self.data.append((real, output, loss))

    def __str__(self):
        report = list()
        # report.append("\n[total]")
        # report.append(f"\n{str(self.total_eval)}")
        for i in range(self.category_len):
            report.append("[{}]".format(self.category_names[i]))
            report.append(f"\n{str(self.get_evaluation(i))}")
        return "".join(report)


class ModelEvaluation:
    def __init__(self, detailed_str=False):
        self.detailed_str = detailed_str
        self.tp = 0
        self.fn = 0
        self.fp = 0
        self.tn = 0

    def add_tp(self):
        self.tp += 1

    def add_fn(self):
        self.fn += 1

    def add_fp(self):
        self.fp += 1

    def add_tn(self):
        self.tn += 1

    def clear(self):
        self.tp = 0
        self.fn = 0
        self.fp = 0
        self.tn = 0

    def get_accuracy(self):
        branch = self.tp + self.tn + self.fp + self.fn
        return (self.tp + self.tn) / branch if branch != 0 else 0

    def get_precision(self):
        branch = self.tp + self.fp
        return self.tp / branch if branch != 0 else 0

    def get_recall(self):
        branch = self.tp + self.fn
        return self.tp / branch if branch != 0 else 0

    def get_f1(self):
        precision = self.get_precision()
        recall = self.get_recall()
        branch = precision + recall
        return 2 * ((precision * recall) / branch) if branch != 0 else 0

    def __str__(self):
        report = list()
        if self.detailed_str:
            accuracy = self.get_accuracy()
            precision = self.get_precision()
            recall = self.get_recall()
            branch = precision + recall
            report.append(
                "\tF1:{:.3f} = 2 * ((precision[{:.3f}] * recall[{:.3f}]) / (precision[{:.3f}] + recall[{:.3f}]))".format(
                    self.get_f1(), precision, recall, precision, recall))
            report.append(
                "\n\taccuracy:{:.3f} = (tp[{:.3f}] + tn[{:.3f}]) / (tp[{:.3f}] + tn[{:.3f}] + fp[{:.3f}] + fn[{:.3f}]))".format(
                    self.get_accuracy(), self.tp, self.tn, self.tp, self.tn, self.fp, self.fn))
            report.append(
                "\n\tprecision:{:.3f} = tp[{:.3f}] / (tp[{:.3f}] + fp[{:.3f}]))".format(self.get_precision(), self.tp,
                                                                                      self.tp, self.fp))
            report.append(
                "\n\trecall:{:.3f} = tp[{:.3f}] / (tp[{:.3f}] + fn[{:.3f}]))".format(self.get_recall(), self.tp, self.tp,
                                                                                   self.fn))
            report.append("\n")
        else:
            report.append("F1:{:.3f}".format(self.get_f1()))
            report.append("\taccuracy:{:.3f}".format(self.get_accuracy()))
            report.append("\tprecision:{:.3f}".format(self.get_precision()))
            report.append("\trecall:{:.3f}".format(self.get_recall()))
            report.append("\n")
        return "".join(report)